import networkx

class Tree(networkx.Graph):
    """A data structure for phylogenetic trees.  Internally, it a subclass of a NetworkX
    graph."""

    def leaves(self): 
        """Generate a list of all leaf nodes in the tree.  For our purposes this is defined as
        all nodes with a degree of 1 and also 0.  Nodes that are not adjacent to other nodes
        are included in the list."""
        return [n for n in self.nodes() if self.degree(n) <= 1]

    def __str__(self):
        return "Tree{0} (nodes:{1} edges:{2})".format(hash(self), len(self.nodes()),
                                                      len(self.edges()))
    def __repr__(self): return self.__str__()
    

class Node():
    """NetworkX expects that all nodes are unique and hashable, yet in phylogenetic trees, 
    the labels of nodes may be missing, and are not necessarily unique.  For this reason,
    nodes are wrapped as instances of the Node class, which makes them uniquely hashable,
    allowing nodes with duplicate or missing labels."""
    
    def __init__(self, label=None):
        self.label = label
        
    def __str__(self):
        return "Node{0} ({1})".format(hash(self), self.label)
    def __repr__(self):
        return self.__str__()
    
